from torch.utils.data import Dataset
from PIL import Image
import os

class MyData(Dataset):
    def __init__(self, root_dir, label_dir):
        self.root_dir = root_dir
        self.label_dir = label_dir
        self.path = os.path.join(root_dir, label_dir)
        self.img_list = os.listdir(self.path)

    def __getitem__(self, index):
        img_name = self.img_list[index]
        img_path = os.path.join(self.path, img_name)
        img = Image.open(img_path)
        img_label = self.label_dir
        return img, img_label

    def __len__(self):
         return len(self.img_list)

class MyDataSet2(Dataset):
    def __init__(self, root_dir, img_dir, label_dir):
        self.root_dir = root_dir
        self.img_dir = img_dir
        self.label_dir = label_dir
        self.img_list = os.listdir(os.path.join(root_dir, img_dir))
        self.label_list = os.listdir(os.path.join(root_dir, label_dir))
        self.img_list.sort()
        self.label_list.sort()

    def __getitem__(self, index):
        img = Image.open(os.path.join(self.root_dir, self.img_dir, self.img_list[index]))
        label = ""
        with open(os.path.join(self.root_dir, self.label_dir, self.label_list[index]), "r") as file:
            label = file.read()
        return img, label

    def __len__(self):
        return len(self.img_list)

root_dir = "datas/train"
ant_label = "ants"
bees_label = "bees"

ants_dataset = MyData(root_dir, ant_label)
bees_dataset = MyData(root_dir, bees_label)

train_dataset = ants_dataset + bees_dataset

# img, label = bees_dataset[0]
# img.show()
# print(label)


root_dir2 = "datas/train2"
ants_dataset2 = MyDataSet2(root_dir2, "ants-img", "ants-label")
bees_dataset2 = MyDataSet2(root_dir2, "bees-img", "bees-label")

train_dataset2 = ants_dataset2 + bees_dataset2
img, label = train_dataset2[123]
img.show()
print(label)



